import os

os.environ["TF_CPP_MIN_LOG_LEVEL"] = '2'
import tensorflow as tf
import numpy as np
from tensorflow.python import keras
from tensorflow.python.keras.preprocessing.image import ImageDataGenerator, load_img, img_to_array
from tensorflow.python.keras.applications.vgg16 import VGG16, preprocess_input


class TransferModel:
    def __init__(self):
        # 初始化能够生成训练数据和测试数据的迭代器对象
        # rescale代表了将数据进行归一化操作
        self.train_generator = ImageDataGenerator(rescale=1.0 / 255.0)
        self.test_generator = ImageDataGenerator(rescale=1.0 / 255.0)

        # 获取数据相关参数
        self.train_path = "./data/train"
        self.test_path = "./data/test"
        self.image_size = (224, 224)
        self.batch_size = 32

        # 创建VGG16不包含最后三层全连接层的模型
        # include_top=False，创建的是不包含最后三层全连接层的模型
        self.base_model = VGG16(weights='imagenet', include_top=False)

        # 初始化类别字典
        self.label_dict = {
            '0': "bus",
            '1': "dinosaurs",
            '2': "elephents",
            '3': "flowers",
            '4': "horse",
        }

    def load_data(self):
        """
        获取数据
        :return: 训练数据和测试数据的迭代器
        """
        # 获取训练数据
        # class_mode有三个参数可选：
        # 1. categorical：将数据的label转换为one-hot编码
        # 2. binary：将数据按label对应的数字编号返回成1维的数组
        # 3. sparse：将label按sparse矩阵返回
        train_gen = self.train_generator.flow_from_directory(directory=self.train_path,
                                                             target_size=self.image_size,
                                                             class_mode="binary",
                                                             batch_size=32,
                                                             shuffle=True)
        # 获取测试数据
        test_gen = self.test_generator.flow_from_directory(directory=self.test_path,
                                                           target_size=self.image_size,
                                                           batch_size=self.batch_size,
                                                           class_mode="binary",
                                                           shuffle=True)

        # return (train_data, test_data), 1  # 这种的返回结果必须用(a, b), c接收，否则会报错
        return train_gen, test_gen

    def refine_vgg_model(self):
        """
        微调VGG模型
        :return:
        """
        # x = self.base_model.outputs
        # print(x)  # [<tf.Tensor 'block5_pool/MaxPool:0' shape=(?, ?, ?, 512) dtype=float32>]
        x = self.base_model.outputs[0]  # 等价于self.base_model.output，input一样
        # print(x)  # Tensor("block5_pool/MaxPool:0", shape=(?, ?, ?, 512), dtype=float32)

        # 对VGG基础模型进行全局池化
        # 做迁移学习不需要调整大量的参数时加上全局池化
        x = keras.layers.GlobalAveragePooling2D()(x)
        # print(x)  # Tensor("global_average_pooling2d/Mean:0", shape=(?, 512), dtype=float32)

        # 添加两个全列阶层
        x = keras.layers.Dense(1024, activation=tf.nn.relu)(x)
        y_predict = keras.layers.Dense(5, activation=tf.nn.softmax)(x)

        # 构建新的模型
        model = keras.models.Model(inputs=self.base_model.inputs, outputs=y_predict)

        return model

    def freeze_base_layers(self):
        """
        冻结VGG基础的5个block层
        :return:
        """
        for layer in self.base_model.layers:
            layer.trainable = False
            # print(layer)

    def compile(self, model):
        """
        编译模型，创建优化器，设置损失，计算准确率
        :param model: 模型
        :return:
        """
        model.compile(optimizer=keras.optimizers.Adam(),
                      loss=keras.losses.sparse_categorical_crossentropy,
                      metrics=["accuracy"])

    @staticmethod
    def fit_generator(model, train_gen, test_gen):
        """
        训练模型
        :param model: 要训练的模型
        :param train_gen: 训练集迭代器
        :param test_gen: 测试集迭代器
        :return:
        """
        # 程序不会则定创建ckpt文件夹，必须手动创建
        check = keras.callbacks.ModelCheckpoint(filepath="./ckpt/{epoch:02d}-{acc:.2f}.h5",
                                                monitor="acc",
                                                save_best_only=True,
                                                save_weights_only=True,
                                                model="auto",
                                                period=1)

        # 只有fit_generator中的callbacks能用keras.callbacks.ModelCheckpoint
        # fit_generator传入的数据是迭代器
        model.fit_generator(generator=train_gen, epochs=3, validation_data=test_gen, callbacks=[check])

    def predict(self, model):
        """
        模型的预测
        :param model: 用于预测的模型
        :return:
        """
        # 加载模型
        model.load_weights("./ckpt/transfer_03-0.99.h5")

        # 加载预测图片
        image = load_img("./image/bus.jpg", target_size=self.image_size)

        # 将图片转换为数组
        image = img_to_array(image)

        # 将图片对应的数组转换为四维结构
        image = image.reshape((1, image.shape[0], image.shape[1], image.shape[2]))

        # 归一化处理
        image = preprocess_input(image)
        print(image.shape)  # (1, 224, 224, 3)

        # 进行预测
        y_predictions = model.predict(image)
        print(y_predictions)  # [[1. 0. 0. 0. 0.]]

        # 对预测结果进行解码
        pre_label = np.argmax(y_predictions, axis=1)
        print(pre_label)  # [0]

        # 返回预测类别
        print(self.label_dict[str(pre_label[0])])


if __name__ == '__main__':
    tm = TransferModel()

    # # 获取数据
    # train_gen, test_gen = tm.load_data()
    # # print(train_gen)
    # # print(test_gen)
    # # for data in train_gen:
    # #     print(data[0].shape)  # (32, 224, 224, 3)
    # #     print(data[1].shape)  # categorical: (32, 5), binary: (32,), sparse: (32,)
    # #     # print(data[1])
    # #     break
    # # print(tm.base_model.summary())
    #
    # 微调模型
    model = tm.refine_vgg_model()
    #
    # # 冻结VGG初始层
    # tm.freeze_base_layers()  # 在训练模型的时候使用
    #
    # # 编译模型
    # tm.compile(model)
    #
    # # 训练模型
    # tm.fit_generator(model, train_gen, test_gen)

    # 模型预测
    tm.predict(model)
